//
// Created by Tommy on 15/3/23.
// Copyright (c) 2015 com.koala. All rights reserved.
//

#import "NSObject+KLObserver.h"
#import "NSObject+KLSwizzle.h"
#import "NSObject+KLExecuteOnDealloc.h"
#import "KLLog.h"
#import <objc/runtime.h>




static char KLKVOObserverInfoArrayKey;
static char KLKVOObserverAllowMethodForwardingKey;

static char KLDeallocBlockArrayKey;

typedef void (^__deallcok_block)(id objc);

@interface __KLKVOObserverInfo : NSObject

@property(nonatomic, copy) NSString *keyPath;
@property(nonatomic, assign) void *context;
@property(nonatomic, assign) void *blockKey;

@end
@implementation __KLKVOObserverInfo


@end

@implementation NSObject (KLObserver)

#if KLENABLE_EXCEPTION_SWIZZLE
+ (void)load{
    static dispatch_once_t onceToken;
    dispatch_once(&onceToken, ^{

        [self kl_swizzleInstanceSelector:@selector(addObserver:forKeyPath:options:context:) withNewSelector:KLHookedSel(addObserver:forKeyPath:options:context:)];
        [self kl_swizzleInstanceSelector:@selector(removeObserver:forKeyPath:context:) withNewSelector:KLHookedSel(removeObserver:forKeyPath:context:)];
        [self kl_swizzleInstanceSelector:@selector(removeObserver:forKeyPath:) withNewSelector:KLHookedSel(removeObserver:forKeyPath:)];
        [self kl_swizzleInstanceSelector:NSSelectorFromString(@"dealloc") withNewSelector:KLHookedSel(dealloc)];

    });
}


- (BOOL)allowMethodForwarding
{
    NSNumber *state = objc_getAssociatedObject(self, &KLKVOObserverAllowMethodForwardingKey);
    return [state boolValue];
}

- (void)setAllowMethodForwarding:(BOOL)allowForwarding
{
    objc_setAssociatedObject(self, &KLKVOObserverAllowMethodForwardingKey, @(allowForwarding), OBJC_ASSOCIATION_RETAIN);
}

- (void)KLHookedMethod(dealloc){

    
    NSMutableArray *blocks = [self deallocBlockArray];
    if(blocks){
        [blocks enumerateObjectsUsingBlock:^(__deallcok_block deallocBlock, NSUInteger idx, BOOL *stop) {
            deallocBlock(self);
        }];
    }
    
    
    [self KLHookedMethod(dealloc)];
}

- (NSMutableArray*)deallocBlockArray{
     return (NSMutableArray *)objc_getAssociatedObject(self, &KLDeallocBlockArrayKey);
}

- (void)setDeallocBlockArray:(NSMutableArray*)blockArray{
     objc_setAssociatedObject(self, &KLDeallocBlockArrayKey,blockArray,OBJC_ASSOCIATION_RETAIN);
}


- (void)__kl_setupDeallocBlockWith:(id)observer forKeyPath:(NSString *)keyPath options:(NSKeyValueObservingOptions)options context:(void *)aContext{
    
    
    NSMutableArray *deallocBlocks = [self deallocBlockArray];
    if (!deallocBlocks) {
        deallocBlocks = [NSMutableArray array];
        [self setDeallocBlockArray:deallocBlocks];
    }
    
     __deallcok_block deallocBlock= ^(id obj) {
         @try {
            [obj removeObserver:observer forKeyPath:keyPath context:aContext];
         }
         @catch (NSException *exception) {
             KLLogError(@"%@",exception);
         }
    } ;
    
    [deallocBlocks addObject:[deallocBlock copy]];
}

- (void)KLHookedMethod(addObserver:(id)observer forKeyPath:(NSString *)keyPath options:(NSKeyValueObservingOptions)options context:(void *)aContext)
{
    //! store info into our observer structure
    NSMutableDictionary *registeredKeyPaths = (NSMutableDictionary *)objc_getAssociatedObject(observer, &KLKVOObserverInfoArrayKey);
    if (!registeredKeyPaths) {
        registeredKeyPaths = [NSMutableDictionary dictionary];
        objc_setAssociatedObject(observer, &KLKVOObserverInfoArrayKey, registeredKeyPaths, OBJC_ASSOCIATION_RETAIN);
    }

    NSMutableArray *observerInfos = registeredKeyPaths[keyPath];
    if (!observerInfos) {
        observerInfos = [NSMutableArray array];
        registeredKeyPaths[keyPath] = observerInfos;
    }
    __block __KLKVOObserverInfo *observerInfo = nil;


    observerInfo = [[__KLKVOObserverInfo alloc] init];
    [observerInfos addObject:observerInfo];


    observerInfo.keyPath = keyPath;
    observerInfo.context = aContext;

    //! Add auto remove when observer is going to be deallocated
    __weak typeof(self) weakSelf = self;
    void *key = [observer kl_performBlockOnDealloc:^(id obj){
        __strong typeof(weakSelf) self = weakSelf;
        id strongObserver = obj;
        NSInteger numberOfRemovals = 0;
        if ((numberOfRemovals = [self kl_removeObserver:strongObserver forKeyPath:keyPath context:aContext registeredKeyPaths:registeredKeyPaths])) {
            for (NSInteger i = 0; i < numberOfRemovals; ++i) {
                [self setAllowMethodForwarding:YES];
                [self KLHookedMethod(removeObserver:observer forKeyPath:keyPath context:aContext)];
                [self setAllowMethodForwarding:NO];
            }
        }
    }];

    observerInfo.blockKey = key;
    
    [self __kl_setupDeallocBlockWith:observer forKeyPath:keyPath options:options context:aContext];
  
    //! call originalMethod
    [self KLHookedMethod(addObserver:observer forKeyPath:keyPath options:options context:aContext)];

}


- (void)KLHookedMethod(removeObserver:(id)observer forKeyPath:(NSString *)keyPath)
{
    if ([self allowMethodForwarding]) {
        [self KLHookedMethod(removeObserver:observer forKeyPath:keyPath)];
        return;
    }

    NSMutableDictionary *registeredKeyPaths = (NSMutableDictionary *)objc_getAssociatedObject(observer, &KLKVOObserverInfoArrayKey);
    NSInteger numberOfRemovals = 0;
    
    if ((numberOfRemovals = [self kl_removeObserver:observer forKeyPath:keyPath context:nil registeredKeyPaths:registeredKeyPaths])) {
        for (NSInteger i = 0; i < numberOfRemovals; ++i) {

            [self setAllowMethodForwarding:YES];
            [self KLHookedMethod(removeObserver:observer forKeyPath:keyPath)];
            [self setAllowMethodForwarding:NO];
        }
    }
}

- (void)KLHookedMethod(removeObserver:(id)observer forKeyPath:(NSString *)keyPath context:(void *)context)
{
    if ([self allowMethodForwarding]) {
        [self KLHookedMethod(removeObserver:observer forKeyPath:keyPath context:context)];
        return;
    }

    NSMutableDictionary *registeredKeyPaths = (NSMutableDictionary *)objc_getAssociatedObject(observer, &KLKVOObserverInfoArrayKey);
    NSInteger numberOfRemovals = 0;
    if ([self allowMethodForwarding] || (numberOfRemovals = [self kl_removeObserver:observer forKeyPath:keyPath context:context registeredKeyPaths:registeredKeyPaths])) {
        for (NSInteger i = 0; i < numberOfRemovals; ++i) {

            [self setAllowMethodForwarding:YES];
            [self KLHookedMethod(removeObserver:observer forKeyPath:keyPath context:context)];
            [self setAllowMethodForwarding:NO];
        }
    }

}

- (NSUInteger)kl_removeObserver:(id)observer
                     forKeyPath:(NSString *)keyPath
                        context:(void *)context
             registeredKeyPaths:(NSMutableDictionary *)registeredKeyPaths
{
    __block NSUInteger result = 0;
    if ([keyPath length] <= 0 && context == nil) {
        //! don't need to execute block on dealloc so cleanup
        [registeredKeyPaths enumerateKeysAndObjectsUsingBlock:^void(id key, id obj, BOOL *stop) {
            NSMutableArray *observerInfos = obj;
            [observerInfos enumerateObjectsUsingBlock:^void(id innerObj, NSUInteger idx, BOOL *innerStop) {
                __KLKVOObserverInfo *info = innerObj;
                [observer kl_cancelDeallocBlockWithKey:info.blockKey];
            }];
        }];
        [registeredKeyPaths removeAllObjects];
        return 1;
    } else {
        [registeredKeyPaths enumerateKeysAndObjectsUsingBlock:^void(id key, id obj, BOOL *stop) {
            NSMutableArray *observerInfos = obj;
            NSMutableArray *objectsToRemove = [NSMutableArray array];
            [observerInfos enumerateObjectsUsingBlock:^void(id innerObj, NSUInteger idx, BOOL *innerStop) {
                __KLKVOObserverInfo *info = innerObj;

                if ((!keyPath || [keyPath isEqualToString:info.keyPath]) && (context == info.context)) {
                    //! remove this info
                    [objectsToRemove addObject:innerObj];

                    //! cancel dealloc block
                    [observer kl_cancelDeallocBlockWithKey:info.blockKey];
                }
            }];

            //! remove all collected objects
            if ([objectsToRemove count] > 0) {
                //! multiple registrations should match unregistrations
                result = [objectsToRemove count];
                [observerInfos removeObjectsInArray:objectsToRemove];
            }
        }];
    }

    return result;
}
#endif
@end
